from matplotlib import image
import torch
import torchvision
import torch.nn.functional as F
import torchvision.transforms as transforms
import torch.optim as optim
import matplotlib.pyplot as plt
import numpy as np
from torchvision.datasets import ImageFolder


from torch.utils.tensorboard import SummaryWriter
from torchvision.transforms.transforms import RandomHorizontalFlip, Resize

'''
 Image preprocessing package in Python, Compose is used to integrate muiltiple steps
--图像预处理化的步骤，将所有的步骤全部写到这个里面
'''
myTransforms  = transforms.Compose(
    [transforms.Resize((40,40)),
    # transforms.RandomResizedCrop(37),
    transforms.ToTensor(),  # 将图像归一化并且变成张量
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]) 
    # 第一个(0.5,0.5,0.5)指图像均值，把图像除0.5
    # 第二个(0.5,0.5,0.5)每个图像都减去0.5
    # (image - 0.5) / 0.5

train_dataset = ImageFolder('/home/lyx/Desktop/LYX_python/dataset/tldataset/train',transform=myTransforms)
# train_dataset = torchvision.datasets.CIFAR10(root='./my_own_test', train = True, download=True,transform = myTransforms)
# train: use train file：是否使用训练文件
# download: if you have not downloaded the test file, program will download from web. if you have downloaded, do nothing
# the method is myTransforms：训练方法是自己定义的Compose
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=2)
# tran_dateset: take test_file form：训练的对象
# batch_size: how many samples per batch to load (default:1)：一次训练时候抓取的样本数量.或者一次检测图片的时候抓取的图片数量
# shuffle: wheather reshuffle after every train：每次训练之前，是否要将训练集的数据打乱
# num_workers: how many subprocesses use for data loading：打开多少个线程加载数据集

# test_dataset = torchvision.datasets.CIFAR10(root='./my_own_test', train=False, download=True, transform=myTransforms)
# test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=4, shuffle=True, num_workers=2)

'''
in_channels: Number of channels in the input image
out_channels: Number of channels produced by the convolution
#注意：输出的通道数量代表了你有多少个卷积核，这里面有32个通道有默认的不同的卷积参数(其实也可以自己更改参数)
#我们训练的实质，其实就是训练这一个参数，对于不同类别的图片，参数会有kernel_size区别，这个参数就相当于是过滤器，通过了这一个过滤器之后我们就能快速得到是哪一类别
kernel_size: Size of the convolving kernel：卷积核的大小
stride：步幅，卷积核在图片上移动后遍历每一个像素点，每次移动的步长就是stride
padding：填充，在图像周围填充一圈像素点。
# 输入图像在进行卷积的过程当中，
'''
class  myNetwork(torch.nn.Module): # 定义自己的网络对象
    def  __init__(self) : 
        '''
        初始化函数
        '''
        super(myNetwork, self).__init__() 
        '''
        super(Net, self).init()是指首先找到Net的父类（比如是类NNet，然后把类Net的对象self转换为类NNet的对象，然后“被转换”的类NNet对象调用自己的init函数
        简单理解：子类把父类的__init__()放到自己的__init__()当中，这样子类就有了父类的__init__()的那些东西。
        '''
        self.convl1 = torch.nn.Conv2d(in_channels = 3, out_channels = 32, kernel_size = 5, stride = 1, padding = 2)
        self.pool11 = torch.nn.MaxPool2d(kernel_size = 2, stride = 2)
        '''
        # 初始化卷积层和池化层
        # 池化层的操作和卷积层的操作类似，但是最后是采用取最大值或者最小值来代表卷积之后那片区域的值
        '''
        self.convl2 = torch.nn.Conv2d(in_channels = 32, out_channels = 32, kernel_size = 5, stride = 1, padding = 2)
        self.pool12 = torch.nn.MaxPool2d(kernel_size = 2, stride = 2)
        self.convl3 = torch.nn.Conv2d(in_channels = 32, out_channels = 32, kernel_size = 5, stride = 1, padding = 2)
        self.pool13 = torch.nn.MaxPool2d(kernel_size = 2, stride = 2)
        # 全连接层 
        ''' 全连接层中的32*4*4是因为输入图像的尺寸为32*32(37-5)*3(彩色图片)，经过卷积之后，变成了32*32*32
            而每一次池化，卷积核的大小为2*2，也就是说，每一次池化之后，图片就会除以4
            32*32*32经过3次池化之后，就变成了(32*32*32)/(4*4*4) = 32*4*4
        '''
        self.fc1 = torch.nn.Linear(in_features = 32*5*5, out_features = 64, bias=True)
        self.fc2 = torch.nn.Linear(in_features = 64, out_features = 64, bias=True)
        self.fc3 = torch.nn.Linear(in_features = 64, out_features = 10, bias=True)
        '''
        in_features指的是输入的二维张量的大小，即输入的[batch_size, size]中的size。
        out_features指的是输出的二维张量的大小，即输出的二维张量的形状为[batch_size，output_size]，当然，它也代表了该全连接层的神经元个数。
        实际上全连接层也是做了一个类似于卷积的操作，将所有的特征点映射到了一个值
        '''
    def forward(self,train_data):
        output = self.pool11(torch.nn.functional.relu(self.convl1(train_data)))
        output = self.pool12(torch.nn.functional.relu(self.convl2(output)))
        output = self.pool13(torch.nn.functional.relu(self.convl3(output)))
        output = output.view(-1, 32*5*5)# 动态调整维度上的元素个数，确保元素的总数不变。
        output = torch.nn.functional.relu(self.fc1(output))
        output = torch.nn.functional.relu(self.fc2(output))
        output = self.fc3(output)
        return output
# 开始
myModel = myNetwork() # 建立了一个模型是myModel
myOptimzier = optim.SGD(myModel.parameters(),lr = 0.01, momentum = 0.9) # 神经网络优化器，主要目的是优化神经网络，使训练过程加快
myLoss = torch.nn.CrossEntropyLoss()    # 计算损失函数


for _epoch in range(10): 
    training_loss = 0.0
    '''
    step每次都是增加1,即遍历了整个图片  
    而每次抓取图片的数量为64
    '''
    for _step, data in enumerate(train_loader): # enumerate()函数的作用是通过迭代来遍历一个字符串、列表或字典等，并且为其增加索引，返回值为enumerate类。
        image,label = data # train_loader里面每个数据集对应了一个图片以及标签
        predict_label = myModel.forward(image)  # 经过模型之后的预测值
        loss = myLoss(predict_label, label) # 预测的标签值和实际值的损失值
        myOptimzier.zero_grad()# optimizer.zero_grad()意思是把梯度置零，也就是把loss关于weight的导数变成0.
        loss.backward() # 损失传递后向更新
        myOptimzier.step()  # 更新所有参数
        training_loss = training_loss + loss.item()
        ''''
         loss是张量，loss.item是对其取数值，但注意 item() 只适用于 tensor 只包含一个元素的时候。
        每10次把平均损失值计算出来
        '''
        if _step % 10 == 0:
            print('[iteration - %3d] training loss: %.3f' % (_epoch*len(train_loader) + _step, training_loss/10)) # 训练的总次数，每10次损失的值/10得到平均值
            training_loss = 0.0

'''保存模型参数'''
# torch.save(myModel.state_dict(), 'params.pkl')
torch.save(myModel, 'model2.pkl')
correct = 0
total = 0
'''
关于with,with的作用是对于with后面任务进行事前设置，事后清理。对异常情况进行处理
（１）紧跟with后面的语句被求值后，返回对象的“–enter–()”方法被调用，这个方法的返回值将被赋值给as后面的变量；
（２）当with后面的代码块全部被执行完之后，将调用前面返回对象的“–exit–()”方法。
    比如文件需要打开，用完需要关闭，并且于对文件读取的数据发生异常的时候，进行处理

对于torch.no_grad()的理解,pytorch的张量变量有一个requires_grad参数，设置为ture的话，反向传播的时候tensor会自动求导
这里设置为no_grad。即不用构建计算图，因为我们只是进行模型测试，不需要进行更新参数，使得内存节约
'''
with torch.no_grad():
    for images,labels in train_loader:
        outputs = myModel(images)
        numbers,predicted = torch.max(outputs.data,1)
        # 函数返回两个张量，一个output_data里面每一行最大的值，第二个是每一行最大值的索引
        total += labels.size(0)
        correct += (predicted==labels).sum().item()

print('Testing Accuracy : %d %%' % ( 100 * correct / total))


